# /// script
# dependencies = [
#   "numpy",
#   "torch",
#   "kernels",
# ]
# ///
import torch
from kernels import get_kernel

torch.manual_seed(42)
sage_attention = get_kernel("drbh/sage-attn-test", version=2)

device = "cuda"
B, H, L, D = 1, 8, 256, 64
q = torch.randn(B, H, L, D, dtype=torch.bfloat16, device=device)
k = torch.randn(B, H, L, D, dtype=torch.bfloat16, device=device)
v = torch.randn(B, H, L, D, dtype=torch.bfloat16, device=device)

out = sage_attention.sageattn3_blackwell(q, k, v)
print(f"sageattn output shape: {out.shape}")